
/*
** application layer multicast server
** location: root of multicast-tree
** use UDP && TCP
** gcc -lwsock32 to compile
*/

#include <stdio.h>
#include <stdlib.h>
#include <Winsock2.h>
#include <process.h>
#include <time.h>

#pragma comment(lib, "ws2_32.lib")

#define SERVER_PORT 12500
#define BUF_SIZE    1024
#define WSVERS      MAKEWORD(2, 2)
#define MAX_CHILD   2

struct Node {
    SOCKET sock;
    struct sockaddr_in addr;
} childNode[MAX_CHILD];
int child_count = 0;

int testRecv(int sLen) {
    if (sLen == 0 || sLen == SOCKET_ERROR) {
        return -1;
    } else if (sLen > 0) {
        return 1;
    }
}

void testSend(int sLen) {
    switch (sLen) {
    case SOCKET_ERROR: printf("send() failed: %d\n", WSAGetLastError()); break;
    case 0: printf("send() failed: %d\n", WSAGetLastError()); break;
    default: break;
    }
    return ;
}

unsigned WINAPI multicast(void *arg) {
    SOCKET hsock = *((SOCKET*)arg);
    char wBuf[BUF_SIZE];
    int sLen;
    int clntAddrLen = sizeof(struct sockaddr_in);

    while (true) {
        printf("Please input multicast message: ");
        // input multicast message
        memset(wBuf, 0, sizeof(wBuf));
        fgets(wBuf, BUF_SIZE, stdin);
        // send to all child nodes
        for (int i = 0; i < child_count; ++i) {
            sLen = send(childNode[i].sock, wBuf, strlen(wBuf), 0);
            testSend(sLen);
        }
    }

    return 0;
}


unsigned WINAPI buildtree(void *arg) {
    SOCKET hsock = *((SOCKET*)arg);
    SOCKET clntSock;
    struct sockaddr_in clntAddr;
    int clntAddrLen;
    char rBuf[BUF_SIZE];
    char wBuf[BUF_SIZE];
    int sLen;
    clntAddrLen = sizeof(clntAddr);

    // listen
    if (listen(hsock, 5) == SOCKET_ERROR) {
        printf("listen() failed: %d\n", WSAGetLastError());
    }

    while (true) {
        printf("tree-server listening...\n");

        // accept a join-request socket
        clntSock = accept(hsock, (SOCKADDR*)&clntAddr, &clntAddrLen);
        if (clntSock == INVALID_SOCKET) {
            printf("eccept() failed: %d\n", WSAGetLastError());
        }
        printf("accept a join request from %s:%d \n",
            inet_ntoa(clntAddr.sin_addr), ntohs(clntAddr.sin_port));

        // reply its port
        memset(wBuf, 0, sizeof(wBuf));
        sprintf(wBuf, "%d", ntohs(clntAddr.sin_port));
        sLen = send(clntSock, wBuf, strlen(wBuf), 0);
        testSend(sLen);

        // deal with the request
        closesocket(clntSock);
        if (child_count < MAX_CHILD) {

            // make it be a directed child
            // create socket
            childNode[child_count].sock = socket(PF_INET, SOCK_STREAM, 0);
            if (childNode[child_count].sock == INVALID_SOCKET)
                printf("socket created error.\n");
            // addr info
            childNode[child_count].addr.sin_family = AF_INET;
            childNode[child_count].addr.sin_addr.s_addr = clntAddr.sin_addr.s_addr;
            childNode[child_count].addr.sin_port = clntAddr.sin_port;
            // connect
            if (connect(childNode[child_count].sock, 
                (SOCKADDR*)&childNode[child_count].addr, 
                clntAddrLen) == INVALID_SOCKET) {
                // may be connected to other parent node
                printf("connect() failed: %d\n", WSAGetLastError());

                continue;
            } else {
                // connect successfully
                memset(wBuf, 0, sizeof(wBuf));
                sprintf(wBuf, "success");
                sLen = send(childNode[child_count].sock, wBuf, sizeof(wBuf), 0);
                if (sLen == 0 || sLen == SOCKET_ERROR) {
                    printf("connect() failed: %d\n", WSAGetLastError());
                } else {
                    printf("%s:%d become a child node.\n",
                        inet_ntoa(childNode[child_count].addr.sin_addr), 
                        ntohs(childNode[child_count].addr.sin_port));

                    child_count++;
                }
                continue;
            }

        } else {
                srand(time(0));
                int i = rand() % child_count;

            // transfer to child node(s)
            //for (int i = 0; i < child_count; ++i) {
                // request message
                memset(wBuf, 0, sizeof(wBuf));
                sprintf(wBuf, "new_node");
                sLen = send(childNode[i].sock, wBuf, strlen(wBuf), 0);
                testSend(sLen);
            //}
            //for (int i = 0; i < child_count; ++i) {
                // IP and port
                memset(wBuf, 0, sizeof(wBuf));
                sprintf(wBuf, "%s,%d", inet_ntoa(clntAddr.sin_addr), 
                    ntohs(clntAddr.sin_port));
                sLen = send(childNode[i].sock, wBuf, strlen(wBuf), 0);
                testSend(sLen);
            //}

        }
    }

    return 0;
}


int main(int argc, char *argv[]) {
    HANDLE   hCastThread, hTreeThread;
    WSADATA  wsaData;
    // winsock version
    if (WSAStartup(WSVERS, &wsaData) != 0)
        printf("WSAStartup() error! \n");

    // create TCP socket
    SOCKET hsock = socket(PF_INET, SOCK_STREAM, 0);
    if (hsock == INVALID_SOCKET)
        printf("TCP socket creation error! \n");

    // server address info
    struct sockaddr_in servAddr;
    memset(&servAddr, 0, sizeof(servAddr));
    servAddr.sin_family = AF_INET;
    servAddr.sin_addr.s_addr = htonl(INADDR_ANY);
    servAddr.sin_port = htons(SERVER_PORT);

    // binding
    if (bind(hsock, (SOCKADDR*)&servAddr, sizeof(servAddr)) == SOCKET_ERROR)
        printf("bind() error! \n");

    // start 
    // waiting for join-multicast-tree request
    // multicast to child nodes over application
    hCastThread = (HANDLE)_beginthreadex(NULL, 0, &multicast, (void*)&hsock, 0, NULL);
    hTreeThread = (HANDLE)_beginthreadex(NULL, 0, &buildtree, (void*)&hsock, 0, NULL);

    WaitForSingleObject(hCastThread, INFINITE);
    WaitForSingleObject(hTreeThread, INFINITE);

    closesocket(hsock);
    WSACleanup();
    return 0;
}

